#!/usr/bin/env python
# Copyright © 2010-2013 Kris Maglione <maglione.k@gmail.com>
#
# This work is licensed for reuse under an MIT license. Details are
# given in the LICENSE file included with this file.

from functools import partial
from itertools import chain
import hashlib
import io
import os
import re
import json as simplejson
import sys

DB_PATH = os.environ.get("EXTENSION_TEST_DB",
                         "%(HOME)s/extension-test/checksums.sqlite" % os.environ)

import sqlite3

MAGIC = {
    'BTPInstaller.jsm': (r'(?m)^var extensionID = "([^"]+)"',
                         r'(?m)^var newBTPVersion = "([^"]+)"'),
    'prefman.js': (r'^function (.*?)_PrefManager\b',),
    'script-compiler-overlay.xul': (r'^<\?xml(?:.|\n)*?chrome://([^/]+)/',),
    'script-compiler.js': (
        r'^var (.*?)_gmCompiler\b',
        r'''(?x)
            \s*
            if \s \( \s* greasemonkeyScriptCompiler_gmCompiler\.isGreasemonkeyable\(href\)
            (?:
                    \s* && \s*
                    !? \s*
                    \( \s*
                    (?:
                        / (?: [^\\/] | \\. )+ /
                        
                        \.test\(href\)

                        \s*

                        (?: \|\| | \) )

                        \s*
                    )*
            )*
            ( \s* && \s* true \s* )*
            \) \s* \{ \s*
                    var \s script=greasemonkeyScriptCompiler_gmCompiler\.getUrlContents\( \s*
                            'chrome://greasemonkeyScriptCompiler/content/[^'\\]*' \s*
                    \); \s*
                    greasemonkeyScriptCompiler_gmCompiler\.injectScript\(script, \s href, \s unsafeWin\); \s*
            \} \s*
        '''
    ),
    'xmlhttprequester.js': (r'^function (.*?)_xmlhttpRequester\b',),
};
for filename, regex in MAGIC.items():
    MAGIC[filename] = tuple(map(re.compile, regex))

hack = set()
def checksum(file=None, name=None, algorithm='sha1', data=None, munger=None):
    try:
        sum = hashlib.new(algorithm)

        try:
            if data is None:
                data = re.sub('^\ufeff', '', open(file, 'rU').read())

            if munger:
                data = munger(data)

            if name is None:
                name = os.path.basename(file)

            if name in MAGIC:
                match = MAGIC[name][0].search(data)
                if match:
                    data = re.sub(r'(\b|_)%s(?=\b|_)' % match.group(1),
                                  lambda m: m.group(1) + "greasemonkeyScriptCompiler",
                                  data)
                    for regex in MAGIC[name][1:]:
                        data = regex.sub("", data)
                    data = re.sub(r'\s*$', "", data)

            if file and any(s in data
                            for s in ('\n;/*', ';/*\n', ';/*!', '/*! jQuery UI')):
                hack.add(file);

            sum.update(data.encode('UTF-8'))
        except UnicodeDecodeError:
            sum.update(open(file, 'rb').read())

        return sum.hexdigest();

    except:
        import traceback
        traceback.print_exc(file=sys.stderr)
        return ''

def quote(str):
    return "'%s'" % str.replace("'", "''")

def lookup(sums):
    group = lambda l: ', '.join(map(quote, l))

    query = """
        SELECT
            sha1sum, title
        FROM
            checksums
        JOIN
            groups USING (groupid)
        WHERE
            sha1sum IN (%s) AND groupname IN (%s);
    """ % (group(sums), group(databases))

    return dict(database.execute(query).fetchall())

exists = os.path.exists(DB_PATH)
database = sqlite3.connect(DB_PATH)
if not exists:
    database.execute("""
            CREATE TABLE checksums (sha1sum CHAR(40) PRIMARY KEY,
                                    sha256sum CHAR(64),
                                    title VARCHAR,
                                    groupid INTEGER,
                                    FOREIGN KEY(groupid) REFERENCES groups(groupid),
                                    PRIMARY KEY(sha1sum, groupid));
    """)
    database.execute('CREATE TABLE groups    (groupid INTEGER PRIMARY KEY AUTOINCREMENT, groupname VARCHAR(28));')
    database.execute('CREATE UNIQUE INDEX checksums1_index ON checksums (sha1sum, groupid);')
    database.execute('CREATE UNIQUE INDEX checksums256_index ON checksums (sha256sum, groupid);')

def createdb(name):
    database.execute('INSERT OR IGNORE INTO groups (groupname) VALUES (?);', [name])
databases = ['checksums']
list(map(createdb, databases))

def usage(ret=1):
    print(""" Checksum validator
 ©2011 Kris Maglione

 Usage: checksums [<command> <args>] ...

 Commands:

    db <database>?       Adds a database to the search list. New checksums are
                         added to the last database specified

    add <title> <file>*  Adds the given files to the current database with the
                         given title.

    check <file>*        Checks the given files against the checksums database
                         and reports the results. Return value is the number of
                         unmatched files.
          
    lookup <file>?       Checks the contents read from the standard input
                         against the checksums database. If <file> is given,
                         the contents are processed as if they were read from
                         a file named <file>.

    help                 Display this message.
""", file=sys.stderr)

    sys.exit(ret)

deheaderify = lambda x: re.sub(r'^\ufeff?/\*.*?\*/\s*', '', x, flags=re.S)

is_jquery_ui = False

json = False;
argv = sys.argv[1:]
while argv:
    cmd = argv.pop(0)
    if cmd == "json":
        json = True
    elif cmd == "db":
        if not argv:
            print('Groups: ' + '\n        '.join(sorted(r[0] for r in database.execute('SELECT groupname FROM groups'))))
            sys.exit(0)

        databases.append(argv.pop(0))
        createdb(databases[-1])

    elif cmd == "jquery-ui":
        is_jquery_ui = True

    elif cmd == "add":
        title      = argv[0]
        filenames  = list(filter(os.path.isfile, argv[1:]))
        sha1sums   = list(map(checksum, filenames))
        sha256sums = list(map(partial(checksum, algorithm='sha256'),
                              filenames))

        titles     = ['%s: %s' % (title, name) for name in filenames] # if len(filenames) > 1 else [title]

        if is_jquery_ui:
            sha1sums  .extend(map(partial(checksum, munger=deheaderify),
                                  filenames))
            sha256sums.extend(map(partial(checksum, algorithm='sha256', munger=deheaderify),
                                  filenames))
            titles    .extend(titles)

        group      = database.execute('SELECT groupid FROM groups WHERE groupname = ?;', [databases[-1]]).fetchone()[0]

        database.executemany('''
                INSERT OR IGNORE
                INTO
                    checksums (sha1sum, sha256sum, title, groupid)
                VALUES
                    (?, ?, ?, %d)
            ''' % group,
            list(zip(sha1sums, sha256sums, titles)))
        database.commit()
        argv = []

    elif cmd == "check":
        filenames = list(filter(os.path.isfile, argv))
        checksums = list(map(checksum, filenames))
        prec      = max(len(n) for n in filenames)
        sums      = dict(zip(filenames, checksums))

        existing  = lookup(checksums)
        unknown   = set(filename
                        for filename, sha1 in sums.items()
                        if sha1 not in existing)

        # Check for amalgamated gunk.
        # TODO: Make work with `checksums lookup`
        pat = re.compile(r'(?<=\n);(?=/\*|$)|;(?=/\*[!\n]|$)')
        jquery_pat = re.compile(r'(?<=\)\(jQuery\));\n?(?=\(function\()')

        for filename in list(unknown):
            if filename in hack:
                try:
                    data = re.sub('\ufeff', '', open(filename, 'rU').read())

                    if data.startswith('/*! jQuery UI '):
                        data = deheaderify(data)[:-1]

                        parts = [(checksum(data=part.strip() + ';', name=filename),)
                                 for part in jquery_pat.split(data)
                                 if part]
                    else:
                        parts = [(checksum(data=part.strip(), name=filename),
                                  checksum(data=part.strip() + '\n', name=filename))
                                 for part in pat.split(data)
                                 if part]

                    known = lookup(chain(*parts))
                    if all(any(p in known for p in part)
                           for part in parts):
                        unknown.remove(filename)
                        existing[sums[filename]] = '\n\t'.join(known.values())
                except UnicodeDecodeError:
                    pass

        if json:
            res = dict(known={})
            for file, sha1 in sums.items():
                if sha1 in existing:
                    res['known'][file] = existing[sha1];

            res['unknown'] = list(unknown)

            simplejson.dump(res, sys.stdout)

            sys.exit(len(filenames) - len(existing))
        else:
            sums      = dict(zip(checksums, filenames))

            lines = []
            for sha1, title in existing.items():
                lines.append('%-*s %s' % (prec, sums[sha1], title))
            print('\n'.join(sorted(lines)))

            res = len(filenames) - len(existing)
            if res:
                print('\nUnknown:')
                lines = ['    %s' % filename for filename in unknown]
                print('\n'.join(sorted(lines)))

                sys.exit(res)

        argv = []

    elif cmd == "lookup":
        result = list(lookup([checksum('/dev/stdin',
                                       argv and argv.pop(0))]).values())
        if result:
            print(result[0])
        else:
            sys.exit(1)

    elif cmd == "help":
        usage(0)
    else:
        usage(1)

# vim:se sts=4 sw=4 et ft=python:
